import sys
import os
import glob
import shutil
import json
import cv2

def parse_number(json_path):
    with open(json_path, 'r') as fin:
        labels = json.load(fin)
        number_labels = []
        for obj in labels['shapes']:
            label = obj['label'] 
            point = obj['points']
            left_top = point[0]
            right_bottom = point[1]
            left, top = map(float,left_top)
            right, bottom = map(float, right_bottom)

            left, top, right, bottom = int(left), int(top), int(right), int(bottom)
            number_id = obj['group_id']
            number_labels.append((number_id, left, top, right ,bottom))
    return number_labels

def parse_json(json_path):
    with open(json_path, 'r') as fin:
        label_obj = json.load(fin)
        img_h = label_obj['imageHeight']
        img_w = label_obj['imageWidth']
        labels = []
        for obj in label_obj['shapes']:
            label = obj['label']
            point = obj['points']
            left_top = point[0]
            right_bottom = point[1]
            left, top = map(float,left_top)
            right, bottom = map(float, right_bottom)
            w = right - left
            h = bottom - top
            x_center = (left + w/2)/img_w
            y_center = (top + h/2)/img_h
            w = w/img_w
            h = h/img_h
            labels.append(' '.join(map(str, [label, x_center, y_center, w, h])))
        return labels 

def parse_number_json(json_path):
    with open(json_path, 'r') as fin:
        label_obj = json.load(fin)
        img_h = label_obj['imageHeight']
        img_w = label_obj['imageWidth']
        labels = []
        for obj in label_obj['shapes']:
            label = obj['label']
            point = obj['points']
            left_top = point[0]
            right_bottom = point[1]
            left, top = map(float,left_top)
            right, bottom = map(float, right_bottom)
            w = right - left
            h = bottom - top
            x_center = (left + w/2)/img_w
            y_center = (top + h/2)/img_h
            w = w/img_w
            h = h/img_h
            if label == 'number':
                labels.append(' '.join(map(str, [0, x_center, y_center, w, h])))
            else:
                print("unsportted label", label)
    return labels


def parse_soccer_label(json_path, img_path):
    with open(json_path, 'r') as fin:
        ball_obj = json.load(fin)
        img_h = ball_obj['imageHeight']
        img_w = ball_obj['imageWidth']
        labels = []
        img = cv2.imread(img_path)
        for obj in ball_obj['shapes']:
            label = obj['label']
            point = obj['points']
            left_top = point[0]
            right_bottom = point[1]

            left, top = map(float , left_top)
            right, bottom = map(float, right_bottom)
            img_center = int(img_w/2)
            img_w_new = img_w/2
            if left< img_center:
                img_out = img[:, :img_center, :]
            else:
                img_out = img[:, img_center:, :] 
                left = left-img_center
                right = right-img_center

            w = right - left
            h = bottom - top
            x_center = (left + w/2)/img_w_new
            y_center = (top + h/2)/img_h
            w = w/img_w_new
            h = h/img_h

            if label == 'soccer':
                labels.append(' '.join(map(str, [0, x_center, y_center, w, h])))
            else:
                print("unsupported labels")
    return  img_out, labels
           
    

def parse_label_json(json_path):
    with open(json_path, 'r') as fin:
        label_obj = json.load(fin)
        img_h = label_obj['imageHeight']
        img_w = label_obj['imageWidth']
        labels = []
        for obj in label_obj['shapes']:
            label = obj['label']
            point = obj['points']
            left_top = point[0]
            right_bottom = point[1]
            left, top = map(float,left_top)
            right, bottom = map(float, right_bottom)
            w = right - left
            h = bottom - top
            x_center = (left + w/2)/img_w
            y_center = (top + h/2)/img_h
            w = w/img_w
            h = h/img_h
            #if label == 'number':
            #    labels.append(' '.join(map(str, [0, x_center, y_center, w, h])))

            if label == 'basketball':
                labels.append(' '.join(map(str, [0, x_center, y_center, w, h])))
            #elif label == 'rim':
            #    labels.append(' '.join(map(str, [1, x_center, y_center, w, h])))
            #elif label == 'net':
            #    labels.append(' '.join(map(str, [2, x_center, y_center, w, h])))
            else:
                print("unsportted label", label)
                
        return labels 

    

def custom_ball():
    img_folder = '/mnt/data/huangying/Codes/auto-segment/datas/out_video/crop_ball_playground_imgs/'
    dst_folder = '/mnt/data/huangying/datas/basket_ball_v4/train/'

    for idx, img_path in enumerate(glob.glob(img_folder+"*.jpg")):
        #print(img_path)
        label_json = img_path[:-3]+'json'
        #print(label_json)
        if not os.path.exists(label_json):
            print("== ignore unlabel ", label_json)
            continue
        labels = parse_label_json(label_json)
        img_name = os.path.basename(img_path)
        dst_img = os.path.join(dst_folder, 'images', img_name)
        shutil.copy(img_path, dst_img)
        dst_label_path = os.path.join(dst_folder, 'labels', img_name[:-3] + 'txt')
        write_label_out(dst_label_path, labels)

def parse_label_file(label_path):
    has_ball_label = False
    valid_labels = []
    with open(label_path, 'r') as fin:
        for line in fin.readlines():
            labels = line.split(' ')
            label = labels[0]
            if label == '32': # coco ball label 32
                valid_labels.append(line.strip())
                has_ball_label = True
    return has_ball_label, valid_labels

def parse_coco_person(label_path):
    has_valid_label = False
    valid_labels = []
    with open(label_path, 'r') as fin:
        for line in fin.readlines():
            items = line.strip().split(' ')
            label = items[0]
            if label == '0': # coco person label 0
                valid_labels.append(line.strip())
                has_valid_label = True
            if label == '32':
                valid_labels.append(line.strip())
                has_valid_label = True
    return has_valid_label, valid_labels
            


def write_ball_label_out(label_path, valid_balls):
    with open(label_path, 'w') as fout:
        for ball in valid_balls:
            items = ball.split(' ')
            items[0] = '0' 
            line = ' '.join(items)
            fout.write(line+'\n')

def write_coco_person_ball_out(label_path, valid_labels):
    with open(label_path, 'w') as fout:
        for label in valid_labels:
            items = label.split(" ")
            if items[0] == '0':
                items[0] = '1'
            elif items[0] == '32':
                items[0] = '0'
            else:
                print("== unsupported label ", items[0], label_path)
            line = ' '.join(items)
            fout.write(line+'\n')
            

def write_rim_label_out(label_path, valid_labels):
    with open(label_path, 'w') as fout:
        for label in valid_labels:
            fout.write(label+'\n')

def coco_ball():
    label_folder = '/home/avs/datasets/coco2017/coco/labels/train2017' 
    img_folder = '/home/avs/datasets/coco2017/train2017'
    dst_folder = '/mnt/data/huangying/datas/soccer_datas/soccer_only/coco'
    for label_file in os.listdir(label_folder):
        label_path = os.path.join(label_folder, label_file)
        has_ball, valid_balls = parse_label_file(label_path)
        if has_ball:
            img_name = label_file[:-3] + 'jpg'
            dst_img = os.path.join(dst_folder, 'images', img_name)
            src_img = os.path.join(img_folder, img_name)
            shutil.copy(src_img, dst_img)
            print("there are ", len(valid_balls), ' ball in ', label_path)
            dst_label_path = os.path.join(dst_folder, 'labels', label_file)
            write_ball_label_out(dst_label_path, valid_balls)
            

def coco_person_ball():
    label_folder = '/home/avs/datasets/coco2017/coco/labels/train2017' 
    img_folder = '/home/avs/datasets/coco2017/train2017'
    dst_folder = '/mnt/data/huangying/datas/coco_person_ball'
    for label_file in os.listdir(label_folder):
        label_path = os.path.join(label_folder, label_file)
        has_label, valid_labels = parse_coco_person(label_path)
        if has_label:
            img_name = label_file[:-3] + 'jpg'
            dst_img = os.path.join(dst_folder, 'images', img_name)
            src_img = os.path.join(img_folder, img_name)
            shutil.copy(src_img, dst_img)
            print("there are ", len(valid_labels), ' person ball in ', label_path)
            dst_label_path = os.path.join(dst_folder, 'labels', label_file)
            write_coco_person_ball_out(dst_label_path, valid_labels)



def rim_basket_ball():
    #input_folder = '/mnt/data/huangying/datas/auto-segment/20230716142423_imgs/'
    #dst_folder = '/mnt/data/huangying/datas/auto-segment/basket_rim_net_v2/'
    #suffix = 'jpeg'
    input_folder = '/mnt/data/huangying/datas/from_hbj/8k_images/'
    dst_folder = '/mnt/data/huangying/datas/cba_jerseynumber_v2'
    suffix = 'jpg'
    for idx, json_path in enumerate(glob.glob(input_folder+"*.json")):
        img_path = json_path[:-len('json')] + suffix
        if idx % 10 == 0: # val
            subset = 'val'
        else:
            subset = 'train'
        #labels = parse_label_json(json_path)       
        labels = parse_number_json(json_path)
        image_name = os.path.basename(img_path)
        dst_img_folder = os.path.join(dst_folder, subset, 'images')
        if not os.path.exists(dst_img_folder):
            os.makedirs(dst_img_folder)
        dst_img_path = os.path.join(dst_img_folder, image_name)
        label_name = image_name[:-len(suffix)] + 'txt'
        dst_label_folder = os.path.join(dst_folder, subset, 'labels')
        if not os.path.exists(dst_label_folder):
            os.makedirs(dst_label_folder)
        label_path = os.path.join(dst_label_folder, label_name)
        shutil.copy(img_path, dst_img_path)
        write_rim_label_out(label_path, labels)

def cba_jersey():
    data_root = '/mnt/data/huangying/datas/cba_basketball_0810_JerseyNumber_train/'
    dst_folder = '/mnt/data/huangying/datas/cba_jerseynumber'
    for idx,label_json in enumerate(glob.glob(os.path.join(data_root, 'json_labels/*.json'))):
        img_name = os.path.basename(label_json)[:-4]+'png'
        img_path = os.path.join(data_root, 'images', img_name)
        #print(img_path)
        if idx % 10 == 0: # val
            subset = 'val'
        else:
            subset = 'train'
        labels = parse_number_json(label_json)
        dst_img_folder = os.path.join(dst_folder, subset, 'images')
        if not os.path.exists(dst_img_folder):
            os.makedirs(dst_img_folder)
        dst_img_path = os.path.join(dst_img_folder, img_name)
        shutil.copy(img_path, dst_img_path)

        dst_label_folder = os.path.join(dst_folder, subset, 'labels')
        if not os.path.exists(dst_label_folder):
            os.makedirs(dst_label_folder)
        label_name = img_name[:-3]+'txt'
        #print(label_name)
        label_path = os.path.join(dst_label_folder, label_name)
        write_rim_label_out(label_path, labels)

def copy_ball_label():
    label_folder = '/mnt/data/huangying/datas/all_balls/train/labels'
    dst_folder = '/mnt/data/huangying/datas/basket_ball_v4/train/labels'
    img_folder = '/mnt/data/huangying/datas/basket_ball_v4/train/images'
    for img in os.listdir(img_folder):
        img_name = img.split('.')[0]
        label_name = "{}.txt".format(img_name)
        label_path = os.path.join(label_folder, label_name)
        dst_path = os.path.join(dst_folder, label_name)
        shutil.copy(label_path, dst_path)

def read_ball_in(label_path):
    ball_lines = []
    with open(label_path, 'r') as fin:
        for line in fin.readlines():
            items = line.strip().split(" ")
            if(items[0] == '0'):
                ball_lines.append(line.strip())
    return ball_lines
                

def filter_ball():
    label_folder = '/mnt/data/huangying/datas/basket_ball_v4/val/labels/'
    for label_file in os.listdir(label_folder):
        label_path = os.path.join(label_folder, label_file)
        print("== label_path ", label_path)
        ball_labels = read_ball_in(label_path)   
        write_rim_label_out(label_path, ball_labels)

def read_label_in(label_path):
    with open(label_path, 'r') as fin:
        return fin.readlines()

def read_lines_in(label_path):
    lines = []
    with open(label_path, 'r') as fin:
        for line in fin.readlines():
            lines.append(line.strip())
    return lines
        

def write_label_out(out_path, label_lines):
    with open(out_path, 'w') as fout:
        for line in label_lines:
            fout.write("{}\n".format(line))

def merge_person_ball():
    person_label_folder = '/mnt/data/huangying/Codes/ByteTrack/YOLOX_outputs/yolox_s_mix_det/detect_boxes_val'
    ball_label_folder = '/mnt/data/huangying/datas/basket_ball_v4/val/labels'
    dst_folder = '/mnt/data/huangying/datas/ball_person_v1/val/labels'
    for file_name in os.listdir(person_label_folder):
        if file_name.endswith('txt'):
            person_labels = read_label_in(os.path.join(person_label_folder, file_name))
            ball_label_path = os.path.join(ball_label_folder, file_name)
            if os.path.exists(ball_label_path):
                ball_labels = read_label_in(ball_label_path)
                person_labels.extend(ball_labels)
            out_path = os.path.join(dst_folder, file_name)
            write_label_out(out_path, person_labels)
   
def filter_ball():
    all_coco_ball_path = '/mnt/data/huangying/datas/all_balls/train/all_coco_ball.txt'
    basket_ball_path = '/mnt/data/huangying/datas/basket_ball_v4/train/basket_ball.txt'
    all_coco_ball = read_label_in(all_coco_ball_path)
    basket_ball = read_label_in(basket_ball_path) 
    diff_ball = []
    for ball in all_coco_ball:
        if ball not in basket_ball:
            diff_ball.append(ball)
    write_label_out('./remove_ball.txt', diff_ball)

def clean_ball():
    remove_ball_path = '/mnt/data/huangying/datas/remove_ball.txt'
    
    train_folder = "/mnt/data/huangying/datas/ball_person_v2/train/images"

    remove_balls = read_lines_in(remove_ball_path)
    for fname in os.listdir(train_folder):
        if fname in remove_balls:
            print("== remove file ", fname)
            path = os.path.join(train_folder, fname)
            os.remove(path)

def crop_soccer_img(json_path, img_path):
    labels = []
    with open(json_path, 'r') as fin:
        print("== load json ", json_path)
        ball_obj = json.load(fin)
        img_h = ball_obj['imageHeight']
        img_w = ball_obj['imageWidth']
        img = cv2.imread(img_path)
        img_h, img_w, _ = img.shape

        obj = ball_obj['shapes'][0]
        label = obj['label']
        point = obj['points']
        left_top = point[0]
        right_bottom = point[1]

        left, top = map(float, left_top)
        left, top = int(left), int(top)
        right, bottom = map(float, right_bottom)
        w = right-left
        h = bottom-top

        x_left = max(left-320, 0)
        x_top = max(top-320, 0)

        x_right = min(img_w, left+320)
        x_bottom = min(img_h, top+320)
        if(x_left == 0):
            x_right = 640
        if(x_top == 0):
            x_bottom = 640
        if(x_right == img_w):
            x_left = img_w-640
        if(x_bottom == img_h):
            x_top = img_h-640

        img_out = img[x_top:x_bottom, x_left:x_right, :]             
        img_h_new, img_w_new, _ = img_out.shape

        x_center = (left - x_left + w/2)/img_w_new
        y_center = (top - x_top + h/2)/img_h_new

        print("x_center ", x_center, " y_center ", y_center)
        w = w/img_w_new
        h = h/img_h_new

        if label == 'soccer':
            labels.append(' '.join(map(str, [0, x_center, y_center, w, h])))
        else:
            print("unsupported labels")
    return img_out, labels

    

def crop_soccer_data():
    src_folder = '/mnt/data/huangying/datas/soccer_datas/taoyuan'
    dst_folder = '/mnt/data/huangying/datas/soccer_datas/soccer_crop_origin/train'
    for idx, f in enumerate(os.listdir(src_folder)):
        if f.endswith('json'):
            file_name = f[:-4]
            json_path = os.path.join(src_folder, f)
            img_path = os.path.join(src_folder, file_name+'jpg')
            img_out, labels = crop_soccer_img(json_path, img_path)
            dst_img_path = os.path.join(dst_folder, 'images', file_name+'jpg')
            dst_label_path = os.path.join(dst_folder, 'labels', file_name+'txt')
            cv2.imwrite(dst_img_path, img_out)
            write_label_out(dst_label_path, labels)

def prepare_soccer_data():
    src_folder = '/mnt/data/huangying/datas/soccer_datas/taoyuan'
    dst_folder = '/mnt/data/huangying/datas/soccer_datas/soccer_only/train'
    for idx, f in enumerate(os.listdir(src_folder)):
        if f.endswith('json'):
            file_name = f[:-4]
            json_path = os.path.join(src_folder, f)
            img_path = os.path.join(src_folder, file_name+'jpg')
            img_out, labels = parse_soccer_label(json_path, img_path)
            dst_img_path = os.path.join(dst_folder, 'images', file_name+'jpg')
            dst_label_path = os.path.join(dst_folder, 'labels', file_name+'txt')
            cv2.imwrite(dst_img_path, img_out)
            write_label_out(dst_label_path, labels)
       
def filter_empty_line():
    label_folder = "/mnt/data/huangying/datas/soccer_datas/soccer_only/val/labels"
    for label_file in os.listdir(label_folder):
        label_path = os.path.join(label_folder, label_file)
        out_lines = []
        with open(label_path, 'r') as fin:
            for line in fin.readlines():
                if line.strip() != "":
                    out_lines.append(line.strip())
        with open(label_path, "w") as fout:
            for line in out_lines:
                fout.write("{}\n".format(line)) 

def filter_extrem_small_ball():
    label_root = '/mnt/data/huangying/datas/soccer_datas/soccer_only/train/labels'
    extrem_small_list =[]
    extrem_small_path = '/mnt/data/huangying/datas/soccer_datas/soccer_only/train/extrem_small.txt'
    for label_file in os.listdir(label_root):
        label_path = os.path.join(label_root, label_file)
        with open(label_path, 'r') as fin:
            for line in fin.readlines():
                items = line.strip().split()
                w = float(items[3])
                h = float(items[4])
                size = w*h*640*640
                #print("file {} size {}".format(label_file, size))
                if size <= 10:
                    #print("== extrem small file ", label_file)
                    extrem_small_list.append(label_file)
    with open(extrem_small_path, 'w') as fout:
        for ex_file in extrem_small_list:
            fout.write("{}\n".format(ex_file))

def resize_ball_train():
    origin_folder = '/mnt/data/huangying/datas/soccer_datas/soccer_only/train/images'
    resized_folder = '/mnt/data/huangying/datas/soccer_datas/soccer_only/train/resized_images' 
    for idx, img_file in enumerate(os.listdir(origin_folder)):
        img_path = os.path.join(origin_folder, img_file)
        img = cv2.imread(img_path)
        h, w, _ = img.shape
        new_h = int(h/(w/640))
        new_img = cv2.resize(img, (640, new_h))
        #print(new_img.shape)
        dst_path  = os.path.join(resized_folder, img_file)
        cv2.imwrite(dst_path, new_img)

def clean_small_ball():
    extrem_small_path = '/mnt/data/huangying/datas/soccer_datas/soccer_only/train/extrem_small.txt'
    train_path = '/mnt/data/huangying/datas/soccer_datas/soccer_only/train/train.txt'

    train_clean_path = '/mnt/data/huangying/datas/soccer_datas/soccer_only/train/train_clean.txt'
    valid_lines = []
    small_list = []
    with open(extrem_small_path, 'r') as fin:
        for line in fin.readlines():
            small_list.append(line.strip().split('.')[0])

    with open(train_path, 'r') as fin:
        for line in fin.readlines():
            file_name = line.strip().split('/')[-1].split('.')[0]
            if file_name not in small_list:
                valid_lines.append(line.strip())

    with open(train_clean_path, 'w') as fout:
        for line in valid_lines:
            fout.write("{}\n".format(line))

def split_soccer_train_val():
    train_img_folder = '/mnt/data/huangying/datas/soccer_datas/soccer_only/train/images'
    train_label_folder = '/mnt/data/huangying/datas/soccer_datas/soccer_only/train/labels'

    dst_img_folder = '/mnt/data/huangying/datas/soccer_datas/soccer_only/val/images'
    dst_label_folder = '/mnt/data/huangying/datas/soccer_datas/soccer_only/val/labels'

    for idx, img_name in enumerate(os.listdir(train_img_folder)):
        train_img_path = os.path.join(train_img_folder, img_name) 
        if idx % 10 == 0:
            print("== move ", idx)
            dst_img_path = os.path.join(dst_img_folder, img_name)
            train_label_path = os.path.join(train_label_folder, img_name[:-3]+'txt')
            dst_label_path = os.path.join(dst_label_folder, img_name[:-3]+'txt')
            shutil.move(train_img_path, dst_img_path)
            shutil.move(train_label_path, dst_label_path)

def cba_player_num_crop():
    input_folder = '/mnt/data/huangying/datas/view_014_frames'
    dst_folder = '/mnt/data/huangying/datas/jersey_datas/cba_v14_numbers'
    dst_path = '/mnt/data/huangying/datas/jersey_datas/cba_v14_numbers.lst'
    with open(dst_path, 'w') as fout:
        for f in os.listdir(input_folder):
            if f.endswith('json'):
                fname = f[:-5]
                img_name = fname+'.bmp'
                label_json = os.path.join(input_folder, f)
                img_path = os.path.join(input_folder, img_name)
                labels = parse_number(label_json)
                img = cv2.imread(img_path)
                for idx, label in enumerate(labels):
                    g_id, left, top, right, bottom = label
                    out_name = "{}_{}_{}.jpg".format(fname, idx, g_id)
                    out_path = os.path.join(dst_folder, out_name)
                    out_label = "{}\t{}".format(out_name, g_id)
                    fout.write("{}\n".format(out_label))
                    label_img = img[top:bottom, left:right,:]
                    cv2.imwrite(out_path, label_img)
                
def select_xw_train():
    xw_all_dir = '/mnt/nas/07_自动导播素材及处理后片段/07新屋村足球场/20240725/20240725_imgs'
    dst_dir = '/mnt/data/huangying/datas/soccer_datas/xw_0725'
    for idx,img in enumerate(os.listdir(xw_all_dir)):
        if idx % 10 == 0:
            in_path = os.path.join(xw_all_dir, img)
            dst_path = os.path.join(dst_dir, img)
            shutil.copy(in_path, dst_path)

def parse_number_v3():
    img_folder = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/cba_number_v3'
    dst_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/cba_number_v3_croped'
    for label_file in os.listdir(img_folder):
        if label_file.endswith('json'):
            label_name = label_file.split('/')[-1][:-5]
            print(label_name)
            json_path = os.path.join(img_folder, label_file)
            img_name =  '{}.jpg'.format(label_name)
            img_path = os.path.join(img_folder, img_name)
            dst_img = os.path.join(dst_dir, 'images', img_name)
            dst_label = os.path.join(dst_dir, 'labels', '{}.txt'.format(label_name))
            labels = parse_number_json(json_path)
            write_label_out(dst_label, labels)
            shutil.copy(img_path, dst_img)

def parse_basket_label(json_path):
    with open(json_path, 'r') as fin:
        labels = json.loads(fin.read())
        basket_labels = labels["shapes"]
        basket_points = []
        for basket in basket_labels:
            if basket['label'] == 'basketball':
                points = basket['points']
                #print("== soccer points ", points)
                basket_points.append(points)
        return basket_points

def crop_basket_wh(crop_w, crop_h):
    input_label_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/三人篮球_4K_3_basket/third_basket_labels/'
    input_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/三人篮球_4K_3_basket/third_basket_images/'
    #input_label_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/三人篮球_4K_5_basket/clip_images_labels/'
    #input_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/三人篮球_4K_5_basket/clip_images/'
    dst_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/third_basket_2560/'
    for idx, json_path in enumerate(glob.glob(input_label_dir+"*.json")):
        if idx % 20 == 0:
            print("handle ", idx)
        img_name = json_path.split('/')[-1].split('.')[0]
        img_path = os.path.join(input_dir, "{}.jpg".format(img_name))
        image = cv2.imread(img_path)
        print("== parse json ", json_path)
        try:
            basket_rect_labels = parse_basket_label(json_path)
        except Exception as e:
            print("======== exception parse ", json_path)
            continue
        for b_idx, basket_points in enumerate(basket_rect_labels):
            left, top = basket_points[0]
            right, bottom = basket_points[1]
            basket_w = right-left
            basket_h = bottom-top
            x_idx = int(left/crop_w)
            y_idx = int(top/crop_h)
            print("x_idx ", x_idx, ", y_idx", y_idx, ", left ", left)
            img = image[y_idx*crop_h:(y_idx+1)*crop_h, x_idx*crop_w:(x_idx+1)*crop_w,  :]
            
            img_h, img_w, _ = img.shape
            l_left = (left-x_idx*crop_w+basket_w/2)/img_w
            l_top = (top-y_idx*crop_h+basket_h/2)/img_h
            l_w = basket_w/img_w
            l_h = basket_h/img_h
            label = '0 {} {} {} {}'.format(l_left, l_top, l_w, l_h)
            out_label_path = os.path.join(dst_dir, 'labels', '{}_{}.txt'.format(img_name, b_idx))
            write_label_out(out_label_path, [label])
           
            out_img_path = os.path.join(dst_dir, 'images',  '{}_{}.jpg'.format(img_name, b_idx))
            cv2.imwrite(out_img_path, img)

    
       
def crop_basket_640():
    #input_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/cba_multiview/cba_gd_bj_20241212_org_unpack/detect_result/0_res_imgs/'
    #dst_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/cba_multiview/train_dst_crop_640'
    input_label_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/三人篮球_4K_5_basket/clip_images_labels/'
    input_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/三人篮球_4K_5_basket/clip_images/'
    dst_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/third_basket_air/'
    crop_w = 640
    crop_h = 640
    for idx, json_path in enumerate(glob.glob(input_label_dir+"*.json")):
        if idx % 20 == 0:
            print("handle ", idx)
        img_name = json_path.split('/')[-1].split('.')[0]
        img_path = os.path.join(input_dir, "{}.jpg".format(img_name))
        image = cv2.imread(img_path)
        print("== parse json ", json_path)
        try:
            basket_rect_labels = parse_basket_label(json_path)
        except Exception as e:
            print("======== exception parse ", json_path)
            continue
        for b_idx, basket_points in enumerate(basket_rect_labels):
            left, top = basket_points[0]
            right, bottom = basket_points[1]
            basket_w = right-left
            basket_h = bottom-top
            x_idx = int(left/crop_w)
            y_idx = int(top/crop_h)
            #print("x_idx ", x_idx, "left ", left)
            img = image[y_idx*crop_h:(y_idx+1)*crop_h, x_idx*crop_w:(x_idx+1)*crop_w,  :]
            
            img_h, img_w, _ = img.shape
            l_left = (left-x_idx*crop_w+basket_w/2)/img_w
            l_top = (top-y_idx*crop_h+basket_h/2)/img_h
            l_w = basket_w/img_w
            l_h = basket_h/img_h
            label = '0 {} {} {} {}'.format(l_left, l_top, l_w, l_h)
            out_label_path = os.path.join(dst_dir, 'labels', '{}_{}.txt'.format(img_name, b_idx))
            write_label_out(out_label_path, [label])
           
            out_img_path = os.path.join(dst_dir, 'images',  '{}_{}.jpg'.format(img_name, b_idx))
            cv2.imwrite(out_img_path, img)

def read_lst_in(all_lst):
    with open(all_lst, 'r') as fin:
        return fin.readlines()
            
def split_basket_trainval():
    dst_folder = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/cba_multiview/train_dst_crop_640'
    all_lst = os.path.join(dst_folder, 'all.lst')
    train_lst = []
    val_lst = [] 
    img_list = read_lst_in(all_lst)
    for idx, img_path in enumerate(img_list):
        if idx % 20 == 0:
            val_lst.append(img_path.strip())
        else:
            train_lst.append(img_path.strip())
    train_path = os.path.join(dst_folder, 'train.lst')
    val_path = os.path.join(dst_folder, 'val.lst')
    write_label_out(train_path, train_lst)
    write_label_out(val_path, val_lst)

if __name__ == '__main__':
    #split_basket_trainval()
    #crop_basket_640()
    crop_basket_wh(crop_w=2560, crop_h=2160)

    #parse_number_v3()
    #select_xw_train()
    #cba_player_num_crop()
    #split_soccer_train_val()
    #crop_soccer_data()
    #resize_ball_train()
    #clean_small_ball()
    #filter_extrem_small_ball()
    #filter_empty_line()
    #prepare_soccer_data()
    #clean_ball()
    #filter_ball()
    #coco_person_ball()
    #merge_person_ball()
    #copy_ball_label()
    #coco_ball()
    #custom_ball()
    #rim_basket_ball()
    #cba_jersey()
